﻿Friend Class dbiSqlClient
    Implements iDBInfo

    Private WithEvents _connection As SqlClient.SqlConnection
    Private _connectionName As String
    Private _messages As New Queue(Of PSObject)
    Private _transaction As SqlClient.SqlTransaction

    Public Property CommandTimeout As Integer Implements iDBInfo.CommandTimeout

    Public ReadOnly Property HasMessages As Boolean Implements iDBInfo.HasMessages
        Get
            Return Me._messages.Count > 0
        End Get
    End Property

    Public Function GetMessage() As PSObject Implements iDBInfo.GetMessage
        Return Me._messages.Dequeue
    End Function

    Public Sub ClearMessages() Implements iDBInfo.ClearMessages
        Me._messages.Clear()
    End Sub

    Public ReadOnly Property Connection As IDbConnection Implements iDBInfo.Connection
        Get
            Return Me._connection
        End Get
    End Property

    Public ReadOnly Property ConnectionInfo As PSObject Implements iDBInfo.ConnectionInfo
        Get
            Dim PS As New PSObject
            With PS.Members
                .Add(New PSNoteProperty("ConnectionName", Me.ConnectionName))
                .Add(New PSNoteProperty("ConnectionType", Me.Type.ToString))
                .Add(New PSNoteProperty("ConnectionState", Me.Connection.State))
                .Add(New PSNoteProperty("ConnectionString", Me.Connection.ConnectionString))
                .Add(New PSNoteProperty("ConnectionTimeout", Me.Connection.ConnectionTimeout))
                .Add(New PSNoteProperty("ServerVersion", Me._connection.ServerVersion))
                .Add(New PSNoteProperty("Database", Me._connection.Database))
                .Add(New PSNoteProperty("CommandTimeout", Me.CommandTimeout))
                .Add(New PSNoteProperty("HasTransaction", Me._transaction IsNot Nothing))
            End With
            Return PS
        End Get
    End Property

    Public ReadOnly Property ConnectionName As String Implements iDBInfo.ConnectionName
        Get
            Return Me._connectionName
        End Get
    End Property

    Friend Function CreateMap(dt As DataTable) As List(Of DBInfo.Map) Implements iDBInfo.CreateMap
        Return dt.AsEnumerable.Select(Function(x)
                                          If String.IsNullOrEmpty(x("ColumnName")) Then
                                              Return New DBInfo.Map(x("ColumnOrdinal"), String.Format("Column{0}", x("ColumnOrdinal")), x("DataType"), x("AllowDBNull"))
                                          Else
                                              Return New DBInfo.Map(x("ColumnOrdinal"), x("Columnname"), x("DataType"), x("AllowDBNull"))
                                          End If
                                      End Function).ToList
    End Function

    Public Function GetCommand(query As String, commandTimeout As Integer, Optional parameters As Hashtable = Nothing) As IDbCommand Implements iDBInfo.GetCommand
        If commandTimeout < 0 Then commandTimeout = Me.CommandTimeout
        Dim cmd As New SqlClient.SqlCommand(query, Me._connection, Me._transaction)
        cmd.CommandTimeout = commandTimeout
        If parameters IsNot Nothing Then
            For Each kvp As DictionaryEntry In parameters
                If kvp.Value Is Nothing Then
                    cmd.Parameters.AddWithValue(kvp.Key, DBNull.Value)
                Else
                    cmd.Parameters.AddWithValue(kvp.Key, kvp.Value)
                End If
            Next
        End If
        Return cmd
    End Function

    Public Function GetCommand(query As String, Optional parameters As Hashtable = Nothing) As IDbCommand Implements iDBInfo.GetCommand
        Return GetCommand(query, Me.CommandTimeout, parameters)
    End Function

    Public Function GetDataSet(cmd As IDbCommand) As DataSet Implements iDBInfo.GetDataSet
        Dim ds As New DataSet
        Using da As New SqlClient.SqlDataAdapter(cmd)
            da.Fill(ds)
        End Using
        Return ds
    End Function

    Public Function GetParameter() As IDataParameter Implements iDBInfo.GetParameter
        Return New SqlClient.SqlParameter
    End Function

    Friend ReadOnly Property Type As DBInfo.DBInfoType Implements iDBInfo.Type
        Get
            Return DBInfo.DBInfoType.MSSQL
        End Get
    End Property

    Public Sub New(name As String, dataSource As String, initialCatalog As String, timeout As Integer)
        Me._connectionName = name
        Me.CommandTimeout = timeout

        Dim cb As New SqlClient.SqlConnectionStringBuilder
        With cb
            .DataSource = dataSource
            .InitialCatalog = initialCatalog
            .IntegratedSecurity = True
            .ApplicationName = "PowerShell"
        End With
        Me._connection = New SqlClient.SqlConnection(cb.ToString)

        Call Me.Init(dataSource.Split("\")(0))
    End Sub

    Public Sub New(name As String, dataSource As String, initialCatalog As String, timeout As Integer, userId As String, password As String)
        Me._connectionName = name
        Me.CommandTimeout = timeout

        Dim cb As New SqlClient.SqlConnectionStringBuilder
        With cb
            .DataSource = dataSource
            .InitialCatalog = initialCatalog
            .UserID = userId
            .Password = password
            .ApplicationName = "PowerShell"
        End With
        Me._connection = New SqlClient.SqlConnection(cb.ToString)

        Call Me.Init(dataSource.Split("\")(0))
    End Sub

    Public Sub New(name As String, connStr As String, timeout As Integer)
        Me._connectionName = name
        Me.CommandTimeout = timeout
        Me._connection = New SqlClient.SqlConnection(connStr)

        Call Me.Init(New SqlClient.SqlConnectionStringBuilder(connStr).DataSource.Split("\")(0))
    End Sub

    Private Sub Init(server As String)
        If Not DBInfo.TestServer(server) Then
            Throw New Exception(String.Format("Server '{0}' cannot be reached.", server))
        End If

        Try
            Me._connection.Open()
        Catch ex As Exception
            Me._connection.Dispose()
            Me._connection = Nothing
            Throw
        End Try
    End Sub

    Private Sub _connection_InfoMessage(sender As Object, e As SqlClient.SqlInfoMessageEventArgs) Handles _connection.InfoMessage
        Dim ps As New PSObject
        With ps.Members
            .Add(New PSNoteProperty("Generated", DateTime.Now))
            .Add(New PSNoteProperty("Message", e.Message))
        End With
        Me._messages.Enqueue(ps)
    End Sub

    Public Function BulkLoad(dataReader As IDataReader, destinationTable As String, columnMap As Hashtable, batchSize As Integer, batchTimeout As Integer, notify As Action(Of Integer)) As Long Implements iDBInfo.BulkLoad
        Using bcp As New SqlClient.SqlBulkCopy(Me._connection, SqlClient.SqlBulkCopyOptions.KeepIdentity, Nothing)
            With bcp
                .DestinationTableName = destinationTable
                .BatchSize = batchSize
                .BulkCopyTimeout = batchTimeout

                If columnMap IsNot Nothing AndAlso columnMap.Count > 0 Then
                    For Each de As DictionaryEntry In columnMap
                        .ColumnMappings.Add(de.Key, de.Value)
                    Next
                End If

                If notify IsNot Nothing Then
                    bcp.NotifyAfter = batchSize
                    Dim rowsCopied As SqlClient.SqlRowsCopiedEventHandler = Sub(sender As Object, e As SqlClient.SqlRowsCopiedEventArgs)
                                                                                notify(e.RowsCopied)
                                                                            End Sub
                    AddHandler bcp.SqlRowsCopied, rowsCopied
                End If
            End With

            bcp.WriteToServer(dataReader)
            dataReader.Close()
            dataReader.Dispose()
        End Using
        Using cmd As SqlClient.SqlCommand = Me._connection.CreateCommand
            cmd.CommandText = String.Format("SELECT COUNT(1) FROM {0}", destinationTable)
            BulkLoad = cmd.ExecuteScalar()
        End Using
    End Function

    Public Sub BeginTransaction() Implements iDBInfo.BeginTransaction
        If Me._transaction IsNot Nothing Then Throw New InvalidOperationException("Cannot BEGIN a transaction when one is already begun.")
        Me._transaction = Me._connection.BeginTransaction
    End Sub

    Public Sub CommitTransaction() Implements iDBInfo.CommitTransaction
        If Me._transaction Is Nothing Then Throw New InvalidOperationException("Cannot COMMIT a transaction when one does not exist.")
        Me._transaction.Commit()
        Me._transaction.Dispose()
        Me._transaction = Nothing
    End Sub

    Public Sub RollbackTransaction(Optional silent As Boolean = False) Implements iDBInfo.RollbackTransaction
        If Me._transaction IsNot Nothing Then
            Me._transaction.Rollback()
            Me._transaction.Dispose()
            Me._transaction = Nothing
        Else
            If Not silent Then Throw New InvalidOperationException("Cannot ROLLBACK a transaction when one does not exist.")
        End If
    End Sub

    Public Sub AttachTransaction(command As IDbCommand) Implements iDBInfo.AttachTransaction
        command.Connection = Me._connection
        If Me._transaction IsNot Nothing Then command.Transaction = Me._transaction
    End Sub
End Class
